import os, sys
import os.path as osp
import numpy as np
import cv2
import json

from tool import filesystem
from paddlex import create_model

handwrite_rec_model = create_model(model_name="PP-OCRv5_server_rec")
# pagenum_rec_model = create_model(model_name="en_PP-OCRv4_mobile_rec")


def handwrite_rec_func(image, debug=0):
    output = handwrite_rec_model.predict(input=image, batch_size=1)
    outputs = list(output)
    if debug:
        for res in outputs:
            # res.print()
            # res.save_to_img(save_path="./output/")
            res.save_to_json(save_path="./output/")

    ret_data = {}
    ret_data["rec_text"] = ""
    ret_data["rec_score"] = 0.0
    if len(outputs) == 0:
        return ret_data
    
    ret_data["rec_text"] = outputs[0]["rec_text"]
    ret_data["rec_score"] = outputs[0]["rec_score"]
    return ret_data


pagenum_help_img_path = "resource/12-tky.png"
pagenum_help_img = cv2.imread(pagenum_help_img_path)

def pagenum_rec_func(image, debug=0):
    if len(image.shape) == 2:
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
    resize_img = cv2.resize(image, (image.shape[1], pagenum_help_img.shape[0]))
    concate_img = np.concatenate((pagenum_help_img, resize_img), axis=1)
    # print("resize_img.shape:", resize_img.shape)
    # print("pagenum_help_img.shape:", pagenum_help_img.shape)
    # print("concate_img.shape:", concate_img.shape)
    
    output = handwrite_rec_model.predict(input=concate_img, batch_size=1)
    outputs = list(output)
    if debug:
        for res in outputs:
            # res.print()
            # res.save_to_img(save_path="./output/")
            res.save_to_json(save_path="./output/")

    ret_data = {}
    ret_data["rec_text"] = ""
    ret_data["rec_score"] = 0.0
    if len(outputs) == 0:
        return ret_data
    
    ret_data["rec_text"] = outputs[0]["rec_text"]
    ret_data["rec_score"] = outputs[0]["rec_score"]
    
    # 还原
    text = "18097324826"
    ret_data["rec_text"] = ret_data["rec_text"][len(text):]
    return ret_data

        
if __name__ == "__main__":
    # data_dir = "/mnt/data2/sj/work/code/qinghai/seal-text-recognition/test.stamp"
    # for img_path in filesystem.get_all_filepath(data_dir, [".jpg", ".png"]):
    #     bgr_img = cv2.imread(img_path)
    #     seal_recognition_func(bgr_img)


    file_path = sys.argv[1]
    file_paths = []
    if osp.isfile(file_path):
        file_paths.append(file_path)
    else:
        file_paths = filesystem.get_all_filepath(file_path, [".jpg", ".png"])
    
    for file_path in file_paths:
        bgr_img = cv2.imread(file_path)
            
        # res = pagenum_rec_func(bgr_img, debug=1)
        res = handwrite_rec_func(bgr_img, debug=1)
        print(osp.basename(file_path), res)

"""
{'res': {'input_path': None, 'page_index': None, 'rec_text': '2020年8月16日', 'rec_score': 0.821458637714386}}
"""
